package server

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aler9/gortsplib"
	"github.com/aler9/gortsplib/pkg/h264"
	"github.com/pion/rtp"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
)

type RtspSession struct {
	path string

	stream    *gortsplib.ServerStream
	publisher *gortsplib.ServerSession

	receiver    map[string]*WeSocketWebrtcConnectionPull // 其它流引用本流
	receiverMux sync.RWMutex

	isSample    bool
	sequencer   rtp.Sequencer
	firstPacket bool
}

func NewRtspSession(ctx *gortsplib.ServerHandlerOnAnnounceCtx) *RtspSession {
	for _, track := range ctx.Tracks {
		x := track.MediaDescription()
		_ = x.MediaName.Media // video/audio

		switch track := track.(type) {
		case *gortsplib.TrackH264:
			_ = track.PayloadType
		case *gortsplib.TrackH265:
		case *gortsplib.TrackVP8:
		case *gortsplib.TrackVP9:
		case *gortsplib.TrackOpus:
		case *gortsplib.TrackPCMA:
		case *gortsplib.TrackPCMU:
		case *gortsplib.TrackGeneric:
		case *gortsplib.TrackJPEG:
		case *gortsplib.TrackMPEG2Audio:
		case *gortsplib.TrackMPEG2Video:
		case *gortsplib.TrackMPEG4Audio:
		}
	}

	tis := &RtspSession{
		path:      ctx.Path,
		stream:    gortsplib.NewServerStream(ctx.Tracks),
		publisher: ctx.Session,

		receiver:    map[string]*WeSocketWebrtcConnectionPull{},
		sequencer:   rtp.NewRandomSequencer(),
		firstPacket: true,
	}

	return tis
}

func (tis *RtspSession) Close() {
	if tis.stream != nil {
		_ = tis.stream.Close()
		tis.stream = nil
	}
}

func (tis *RtspSession) addReceiver(receiver *WeSocketWebrtcConnectionPull) {
	tis.receiverMux.Lock()
	tis.receiver[receiver.UID] = receiver
	tis.receiverMux.Unlock()
}

func (tis *RtspSession) delReceiver(receiver *WeSocketWebrtcConnectionPull) {
	tis.receiverMux.Lock()
	delete(tis.receiver, receiver.UID)
	tis.receiverMux.Unlock()
}

func (tis *RtspSession) onPacketRTP(ctx *gortsplib.ServerHandlerOnPacketRTPCtx) {
	if len(tis.receiver) == 0 {
		return
	}

	var receiverConnections []*WeSocketWebrtcConnectionPull
	tis.receiverMux.RLock()
	for _, connection := range tis.receiver {
		receiverConnections = append(receiverConnections, connection)
	}
	tis.receiverMux.RUnlock()

	var writeSample = func() {
		tracker := tis.stream.Tracks()[ctx.TrackID]
		switch tracker := tracker.(type) {
		case *gortsplib.TrackH264:

			if ctx.H264NALUs == nil {
				return
			}

			var (
				annexBNALUStartCode = []byte{0x00, 0x00, 0x00, 0x01}
				packetBuffer        bytes.Buffer
			)

			// 关键帧前面添加sps/pps
			for _, naLU := range ctx.H264NALUs {
				typ := h264.NALUType(naLU[0] & 0x1F)
				if typ == h264.NALUTypeIDR {
					packetBuffer.Write(annexBNALUStartCode)
					packetBuffer.Write(tracker.SafeSPS())
					packetBuffer.Write(annexBNALUStartCode)
					packetBuffer.Write(tracker.SafePPS())
				}

				packetBuffer.Write(annexBNALUStartCode)
				packetBuffer.Write(naLU)
			}

			sample := media.Sample{
				Data:     packetBuffer.Bytes(),
				Duration: time.Millisecond * 40,
			}

			for _, connection := range receiverConnections {
				if err := connection.WriteSampleOfRTSP(webrtc.RTPCodecTypeVideo, sample); err != nil {
					log.Println(err)
				}
			}
		case *gortsplib.TrackOpus:

		case *gortsplib.TrackPCMA:

		}
	}

	var writeRtp = func() {
		h264T := tis.stream.Tracks()[0].(*gortsplib.TrackH264)
		//tis.getFMTP(h264T)

		h := ctx.Packet.Header
		if false {
			tis.firstPacket = false
			const (
				outputStapAHeader = 0x78
			)

			spsNalu := h264T.SafeSPS()
			ppsNalu := h264T.SafePPS()

			spsLen := make([]byte, 2)
			binary.BigEndian.PutUint16(spsLen, uint16(len(spsNalu)))

			ppsLen := make([]byte, 2)
			binary.BigEndian.PutUint16(ppsLen, uint16(len(ppsNalu)))

			stapANalu := []byte{outputStapAHeader}
			stapANalu = append(stapANalu, spsLen...)
			stapANalu = append(stapANalu, spsNalu...)
			stapANalu = append(stapANalu, ppsLen...)
			stapANalu = append(stapANalu, ppsNalu...)

			h.SequenceNumber = tis.sequencer.NextSequenceNumber()
			pkt := &rtp.Packet{
				Header:  h,
				Payload: stapANalu,
			}
			for _, connection := range receiverConnections {
				if err := connection.WriteRtpOfRTSP(pkt); err != nil {
					log.Println(err)
				}
			}

			// rtsp 推流中关键帧前面要有sps/pps帧
		}

		// h.SequenceNumber = tis.sequencer.NextSequenceNumber()
		pkt := &rtp.Packet{
			Header:      h,
			Payload:     ctx.Packet.Payload,
			PaddingSize: ctx.Packet.PaddingSize,
		}

		// rtp
		for _, connection := range receiverConnections {
			if err := connection.WriteRtpOfRTSP(pkt); err != nil {
				log.Println(err)
			}
		}
	}

	if tis.isSample {
		// sample
		writeSample()
	} else {
		// rtp
		writeRtp()
	}
}

// NewTrackerSample 创建接收流
func (tis *RtspSession) NewTrackerSample(kinds []webrtc.RTPCodecType) []*webrtc.TrackLocalStaticSample {
	var rtpTrackerList []*webrtc.TrackLocalStaticSample

	for _, track := range tis.stream.Tracks() {
		x := track.MediaDescription()
		_ = x.MediaName.Media // video/audio

		switch track := track.(type) {
		case *gortsplib.TrackH264:
			_ = track
			if t, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264},
				"video", "pion"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackH265:
		case *gortsplib.TrackVP8:
		case *gortsplib.TrackVP9:
		case *gortsplib.TrackOpus:
			if t, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus},
				"video", "pion"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackPCMA:
			if t, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypePCMA},
				"video", "pion"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackPCMU:
			if t, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypePCMU},
				"video", "pion"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackGeneric:
		case *gortsplib.TrackJPEG:
		case *gortsplib.TrackMPEG2Audio:
		case *gortsplib.TrackMPEG2Video:
		case *gortsplib.TrackMPEG4Audio:
		}
	}

	tis.isSample = true
	return rtpTrackerList
}

// NewTrackerRTP 创建接收流
func (tis *RtspSession) NewTrackerRTP(kinds []webrtc.RTPCodecType) []*webrtc.TrackLocalStaticRTP {

	var rtpTrackerList []*webrtc.TrackLocalStaticRTP

	for _, track := range tis.stream.Tracks() {
		x := track.MediaDescription()
		_ = x.MediaName.Media // video/audio

		switch track := track.(type) {
		case *gortsplib.TrackH264:

			if t, err := webrtc.NewTrackLocalStaticRTP(
				webrtc.RTPCodecCapability{
					MimeType:    webrtc.MimeTypeH264,
					SDPFmtpLine: tis.getFMTP(track),
				},
				"video", "pion2"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackH265:
		case *gortsplib.TrackVP8:
		case *gortsplib.TrackVP9:
		case *gortsplib.TrackOpus:
			if t, err := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus},
				"audio", "pion2"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackPCMA:
			if t, err := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypePCMA},
				"audio", "pion2"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackPCMU:
			if t, err := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypePCMU},
				"audio", "pion2"); err == nil {
				rtpTrackerList = append(rtpTrackerList, t)
			}
		case *gortsplib.TrackGeneric:
		case *gortsplib.TrackJPEG:
		case *gortsplib.TrackMPEG2Audio:
		case *gortsplib.TrackMPEG2Video:
		case *gortsplib.TrackMPEG4Audio:
		}
	}

	tis.isSample = false
	return rtpTrackerList
}

func (tis *RtspSession) getFMTP(t *gortsplib.TrackH264) string {
	typ := strconv.FormatInt(int64(t.PayloadType), 10)
	fmtp := typ + " packetization-mode=1"

	var tmp []string
	if t.SPS != nil {
		tmp = append(tmp, base64.StdEncoding.EncodeToString(t.SPS))
	}
	if t.PPS != nil {
		tmp = append(tmp, base64.StdEncoding.EncodeToString(t.PPS))
	}
	fmtp += ";sprop-parameter-sets=" + strings.Join(tmp, ",")

	if len(t.SPS) >= 4 {
		fmtp += "; profile-level-id=" + strings.ToUpper(hex.EncodeToString(t.SPS[1:4]))
	}

	return fmtp
}
